import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import cv2
import os
import threading

class ImageSubscriber(Node):
    def __init__(self):
        super().__init__('image_subscriber')
        self.get_logger().info("开始订阅 /image_compressed 话题")
        
        # 订阅图像话题
        self.subscription = self.create_subscription(
            Image,
            '/image_compressed',
            self.image_callback,
            10
        )
        
        self.bridge = CvBridge()
        self.is_recording = False
        self.frame_count = 0  # 用于给保存的图片命名
        self.lock = threading.Lock()  # 用于线程锁
        
        # 创建保存图像的文件夹
        if not os.path.exists('frames'):
            os.makedirs('frames')

    def image_callback(self, msg):
        try:
            # 将ROS图像消息转换为OpenCV格式的BGR图像
            cv_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
        except Exception as e:
            self.get_logger().error(f"图像转换失败: {e}")
            return

        with self.lock:  # 确保线程安全地访问 is_recording
            if self.is_recording:
                # 保存当前帧为图片文件
                filename = f'frames/frame_{self.frame_count:04d}.jpeg'
                cv2.imwrite(filename, cv_image)
                self.frame_count += 1
                self.get_logger().info(f"保存帧 {filename}")

    def start_recording(self):
        with self.lock:
            if not self.is_recording:
                self.is_recording = True
                self.frame_count = 0  # 重置帧计数
                self.get_logger().info("开始保存每一帧")

    def stop_recording(self):
        with self.lock:
            if self.is_recording:
                self.is_recording = False
                self.get_logger().info("停止保存帧")


def handle_input(node):
    while True:
        command = input("按'q'开始保存每一帧，按'w'停止保存，按'e'退出: ")
        if command == 'q':
            node.start_recording()
        elif command == 'w':
            node.stop_recording()
        elif command == 'e':
            rclpy.shutdown()
            break

def main(args=None):
    rclpy.init(args=args)
    node = ImageSubscriber()

    # 启动键盘输入线程
    input_thread = threading.Thread(target=handle_input, args=(node,), daemon=True)
    input_thread.start()

    # 主线程负责处理ROS消息
    rclpy.spin(node)

    # 等待输入线程结束
    input_thread.join()

    node.destroy_node()

if __name__ == '__main__':
    main()
